#include <MemoryManagement.h>
#include <MonitorDriver.h>
#include <Asm.h>
#include <Kernel.h>
#include <ProcessManager.h>

extern Address end;
extern Address initialEsp;

MemoryManagement::MemoryManagement()
{
  placementAddress = (Address)&end ; // TODO: change to multiboot->mod_end
  heapInitialised = false;
  currentDirectory = kernelDirectory = NULL;
}

MemoryManagement::~MemoryManagement()
{
}

void MemoryManagement::initialise(u32int memEnd)
{
  
  // Make enough frames to reach 0x00000000 - memEnd.
  u32int memEndPage = memEnd - (memEnd%PAGE_SIZE); // make sure memEnd is on a page boundary.
  nFrames = memEndPage/PAGE_SIZE;

  frames = new BitArray(nFrames);

  // Make a page directory.
  kernelDirectory = new(true/*page align*/) PageDirectory();
  currentDirectory = kernelDirectory;

  // Map some pages in the kernel heap area.
  // Here we call getPage but not allocFrame. This causes PageTables
  // to be created where nessecary. We can't allocate frames yet because
  // they need to be identity mapped first below.
  for (int i = HEAP_START; i < HEAP_END; i += PAGE_SIZE)
  {
    kernelDirectory->getPage(i, true);
  }
  
  // Map some pages in the user heap area.
  // Here we call getPage but not allocFrame. This causes PageTables
  // to be created where nessecary. We can't allocate frames yet because
  // they need to be identity mapped first below.
  for (int i = USER_HEAP_START; i < USER_HEAP_END; i += PAGE_SIZE)
  {
    kernelDirectory->getPage(i, true);
  }
  
  // Identity map from KERNEL_START to placementAddress.
  int i = 0;
  while (i < placementAddress)
  {
    // Kernel code is readable but not writable from userspace.
    allocFrame( kernelDirectory->getPage(i, true) , false, false);
    i += PAGE_SIZE;
  }
// TODO BUG HERE!
  // Now allocate those pages we mapped earlier.
  for (i = HEAP_START; i < HEAP_START+HEAP_INITIAL_SIZE; i += PAGE_SIZE)
  {
    // Heap is readable but not writable from userspace.
    allocFrame( kernelDirectory->getPage(i, true) , false, false);
  }
  
  for (i = USER_HEAP_START; i < USER_HEAP_START+USER_HEAP_INITIAL_SIZE; i += PAGE_SIZE)
  {
    allocFrame( kernelDirectory->getPage(i, true) , false, true );
  }

  // write the page directory.
  writePageDirectory((Address)kernelDirectory->getPhysical());
  enablePaging();

  // Initialise the heaps.
  heap = Heap(HEAP_START, HEAP_START+HEAP_INITIAL_SIZE, 0xCFFFF000 /* see memory map */);
  userHeap = Heap(USER_HEAP_START, USER_HEAP_START+USER_HEAP_INITIAL_SIZE, 0xDFFFF000);
  
  heapInitialised = true;
  
}

void *MemoryManagement::malloc(u32int size, bool pageAlign, u32int *physicalAddress)
{
  ASSERT(heapInitialised);
  void *addr = heap.allocate(size, pageAlign);
  if (physicalAddress)
  {
    Page *page = kernelDirectory->getPage((Address)addr, false);
    *physicalAddress = page->getFrame() + (Address)addr%PAGE_SIZE;
  }
  return addr;
}

void MemoryManagement::free(void *p)
{
  ASSERT(heapInitialised);
  heap.free(p);
}

void *MemoryManagement::umalloc(u32int size)
{
  ASSERT(heapInitialised);
  return userHeap.allocate(size, false);
}

void MemoryManagement::ufree(void *p)
{
  ASSERT(heapInitialised);
  userHeap.free(p);
}

//
// allocFrame -- maps a page to a frame.
//
void MemoryManagement::allocFrame(Page *p, bool isKernel, bool isWriteable)
{
  if (p->getFrame())
  {
    return;
  }
  else
  {
    // TODO: make this more efficient than O(n).
    u32int frameIdx = frames->firstClear();
    if (frameIdx == (u32int)-1)
    {
      panic("Swap in not implemented.");
    }
    
    frames->set(frameIdx);
    
    p->setPresent(true);
    p->setRw(isWriteable);
    p->setUser(!isKernel);
    p->setFrame(frameIdx*PAGE_SIZE);
  }
}

//
// allocFrame -- maps a page to a frame.
//
Address MemoryManagement::allocFrame()
{
    // TODO: make this more efficient than O(n).
    u32int frameIdx = frames->firstClear();
    if (frameIdx == (u32int)-1)
    {
      panic("Swap in not implemented.");
    }

    frames->set(frameIdx);

    return frameIdx*PAGE_SIZE;
}

void MemoryManagement::freeFrame(Page *p)
{
  u32int frame;
  if (!(frame=p->getFrame()))
  {
    return;
  }
  else
  {
    frames->clear(frame/PAGE_SIZE);
    p->setFrame(0x0);
  }
}

void MemoryManagement::freeFrame(Address frame)
{
  frames->clear(frame/PAGE_SIZE);
}

void MemoryManagement::remapStack()
{

  ASSERT(currentDirectory);

  u32int i;
  // Allocate some space for the stack.
  for( i = STACK_START;
       i > (STACK_START-STACK_INITIAL_SIZE);
       i -= PAGE_SIZE)
  {
    // General-purpose stack is in user-mode.
    allocFrame( currentDirectory->getPage(i, true), false );
  }
  
  // Flush the TLB
  Address oldDir = readPageDirectory();
  writePageDirectory(oldDir);
  
  Address oldStackPointer = readStackPointer();
  Address oldBasePointer  = readBasePointer();
  Address offset          = STACK_START - initialEsp;
  Address newStackPointer = oldStackPointer + offset;
  Address newBasePointer  = oldBasePointer  + offset;
  
  // Copy the stack.
  
  Kernel::memcpy((u8int *)newStackPointer, (const u8int *)oldStackPointer, initialEsp-oldStackPointer);
  
  // Backtrace through the original stack, copying new values into
  // the new stack.
  Address origBasePtrAddr = oldBasePointer;
  Address origBasePtrVal;
  
  for(u32int i = STACK_START; i > STACK_START-STACK_INITIAL_SIZE; i -= 4)
  {
    Address tmp = * (Address*)i;
    if (( oldStackPointer < tmp) && (tmp < initialEsp))
    {
      tmp = tmp + offset;
      Address *tmp2 = (Address*)i;
      *tmp2 = tmp;
    }
  }
  
  writeStackPointer(newStackPointer);
  writeBasePointer(newBasePointer);

}

void MemoryManagement::alignPlacementAddress()
{
  if (placementAddress%PAGE_SIZE) // if it needs aligning at all!
  {
    placementAddress += PAGE_SIZE - placementAddress%PAGE_SIZE;
  }
}

void MemoryManagement::allocateRange(Address startAddress, Address size)
{
  Address endAddress = startAddress+size;
  processManager->getProcess()->memoryMap.map(startAddress, endAddress-startAddress, MemoryMap::Local);
  startAddress &= PAGE_MASK;
  endAddress   &= PAGE_MASK;

  for (Address i = startAddress; i <= endAddress; i += PAGE_SIZE)
  {
    allocFrame( currentDirectory->getPage(i, true), false );
  }
  // Flush page directory.
  flushPageDirectory();
}

u32int MemoryManagement::getKernelHeapSize()
{
  return heap.getSize();
}

u32int MemoryManagement::getUserHeapSize()
{
  return userHeap.getSize();
}

void MemoryManagement::checkIntegrity()
{
  if(heapInitialised)
  {
//     heap.checkIntegrity();
//     userHeap.checkIntegrity();
  }
}
